'''※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※
File Name: get_train_data.py
Author: GID5564
Description: 生成训练数据-数字，总共15个文件夹，每个文件夹下对应的各种字体各种倾斜角的字符图片3900个（字符15类×字体2种×角度20个），图片的大小是24×24像素。保存文件夹:./dataset/
Version: 1.0
Created Time: 22/04/24-19:46:57
※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※※'''
  
#!/usr/bin/python
# -*- coding: UTF-8 -*-

from __future__ import print_function
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import os
import shutil
import time


# 文本对应的文件夹，给每一个分类建一个文件
def createDirs(label_dict,root):
    """
    根据给定的标签字典在根目录下创建相关文件夹。
    :param label_dict: 包含标签和对应字符的字典。
    :param root: 根目录。
    """
    for value,char in label_dict.items():
        train_images_dir = root+"/"+str(value)
        if os.path.isdir(train_images_dir):
            shutil.rmtree(train_images_dir)
        os.makedirs(train_images_dir)

# %% 生成图片
def makeImage(label_dict, font_path, width=24, height=24, rotate=0):
    """
    根据给定的标签字典和字体路径生成图片。

    :param label_dict: 包含标签和对应字符的字典。
    :param font_path: 字体文件的路径。
    :param width: 生成图片的宽度。
    :param height: 生成图片的高度。
    :param rotate: 旋转图片的角度。
    """
    for value, char in label_dict.items():
        # 创建一个黑色背景的图片
        img = Image.new("RGB", (width, height), "black")
        draw = ImageDraw.Draw(img)
        # 加载字体，字体大小为图片宽度的90%
        font = ImageFont.truetype(font_path, int(width * 0.9))
        # 获取字体的宽高
        try:
            x_left_bottom,y_left_bottom,x_right_top,y_right_top= draw.textbbox((0, 0), char, font=font)
            font_width=x_right_top-x_left_bottom
            font_height=y_right_top-y_left_bottom
        except ValueError as e:
            print(f"An error occurred while getting the text size: {e}")
            continue
        # 计算文字的x,y坐标，使文字居中
        # ♥注:y坐标偏下，需优化
        x = (width - font_width) / 2
        y = (height - font_height) / 2
        # 在计算出的坐标处绘制文字，设置文字颜色和字体
        draw.text((x, y), char, (255, 255, 255), font=font)
        # 旋转图片
        img = img.rotate(rotate)
        # 生成图片保存的路径，格式为：dataset/编号/img-编号_r-旋转角度_时间戳.png
        time_value = int(round(time.time() * 1000))
        img_path = f"dataset/{value}/img-{value}_r-{rotate}_{time_value}.png"
        # 保存图片
        img.save(img_path)


def main(label_dict,font_dir):
    """
    根据字体创建单字符图片，24x24像素
    :param label_dict: 包含标签和对应字符的字典。
    :param font_dir: 字体文件的路径。
    """
    for font_name in os.listdir(font_dir):
        # 把每种字体都取出来，每种字体都生成一批图片
        path_font_file = os.path.join(font_dir, font_name)
        # 倾斜角度从-10到10度，每个角度都生成一批图片
        for k in range(-10, 10, 1): 
            # 每个字符都生成图片
            makeImage(label_dict, path_font_file, rotate = k)
        
        
        
if __name__ == "__main__" :
    # %% 要生成的文本
    label_dict = {0: '0', 1: '1', 2: '2', 3: '3', 4: '4', 5: '5', 6: '6', 7: '7', 8: '8', 9: '9', 10: '=', 11: '+', 12: '-', 13: '×', 14: '÷'}
    # 图片根目录
    root="./dataset"
    # %% 存放字体的路径
    font_dir = "/home/aidtool/data/fonts"
        
    createDirs(label_dict,root)
    main(label_dict,font_dir)    
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        